import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.tree.ParseTreeWalker;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class 字节码汇编器 extends AssemblerBaseListener {

    public static 字节码汇编器 汇编(String 字节码) {
        var input = CharStreams.fromString(字节码);
        var lexer = new AssemblerLexer(input);// 词法分析器，将输入文本转换为词法符号
        var tokens = new CommonTokenStream(lexer);// 词法符号流
        var parser = new AssemblerParser(tokens);// 语法分析器
        var 结果 = new 字节码汇编器();
        ParseTreeWalker.DEFAULT.walk(结果, parser.program());
        return 结果;
    }
    
    
    public static final int 初始化字节码长度 = 1024;
    public byte[] 字节码 = new byte[初始化字节码长度];
    public int 数据尺寸; // 通过 .global 赋值

    public List<Object> 常量池 = new ArrayList<>();
    private Map<String, 标签符号> 标签 = new HashMap<>();
    private int 当前地址; // 指令地址指针（IP）；用于填写代码

    public 函数符号 入口函数;

    private final Map<String, Integer> 指令映射 = new HashMap<>(); // 指令名称 -> 指令索引

    public 字节码汇编器() {
        for (int i = 1; i < 字节码定义.指令集.length; i++) {
            指令映射.put(字节码定义.指令集[i].名称.toLowerCase(), i); // 指令不区分大小写
        }
    }

    public int 取字节码尺寸(){
        return 当前地址;
    }

    @Override
    public void exitProgram(AssemblerParser.ProgramContext ctx) {
        // 遍历结束的时候检查一下是否有未定义的标签
        for (String 标签名 : 标签.keySet()) {
            var 标签符号 = (标签符号) 标签.get(标签名);
            if (!标签符号.是否已定义) {
                System.err.println("标签: " + 标签名 + " 未定义");
            }
        }
    }

    @Override
    public void enterGlobals(AssemblerParser.GlobalsContext ctx) {
        数据尺寸 = Integer.parseInt(ctx.INT().getText());
    }

    @Override
    public void enterFunctionDeclaration(AssemblerParser.FunctionDeclarationContext ctx) {
        var 函数符号 = new 函数符号(
                ctx.name.getText(),
                当前地址,
                Integer.parseInt(ctx.a.getText()),
                Integer.parseInt(ctx.lo.getText()));

        if (函数符号.名称.equals("main")) {
            入口函数 = 函数符号;
        }

        if (常量池.contains(函数符号)) {
            常量池.set(常量池.indexOf(函数符号), 函数符号); // 总是使用最后定义的函数符号，由于是练习，此处先不处理重定义的情况
        } else {
            存入常量池(函数符号);
        }
    }

    private void 存入常量池(Object 对象) {
        常量池.add(对象);
    }

    private int 取常量索引(Object 对象) {
        if (!常量池.contains(对象)) {
            常量池.add(对象);
            return 常量池.size() - 1;
        }

        return 常量池.indexOf(对象);
    }


    @Override
    public void enterLabel(AssemblerParser.LabelContext ctx) {
        var 标签名 = ctx.ID().getText();
        var 标签符号 = 标签.get(标签名);

        if (标签符号 == null) {
            标签符号 = new 标签符号(标签名, 当前地址);
            标签.put(标签名, 标签符号);
        } else {
            if (标签符号.是否已定义) {
                throw new RuntimeException("[第" + ctx.start.getLine() + "行]标签: " + 标签名 + " 重定义");
            } else {
                标签符号.地址 = 当前地址;
                标签符号.是否已定义 = true;
                标签符号.解析引用(字节码);
            }
        }
    }


    @Override
    public void enterInstr_noArgs(AssemblerParser.Instr_noArgsContext ctx) {
        生成操作码(ctx.ID().getSymbol());
    }

    @Override
    public void enterInstr_oneArgs(AssemblerParser.Instr_oneArgsContext ctx) {
        生成操作码(ctx.ID().getSymbol());
        生成操作数(ctx.operand());
    }

    @Override
    public void enterInstr_twoArgs(AssemblerParser.Instr_twoArgsContext ctx) {
        生成操作码(ctx.ID().getSymbol());
        生成操作数(ctx.a);
        生成操作数(ctx.b);
    }

    @Override
    public void enterInstr_threeArgs(AssemblerParser.Instr_threeArgsContext ctx) {
        生成操作码(ctx.ID().getSymbol());
        生成操作数(ctx.a);
        生成操作数(ctx.b);
        生成操作数(ctx.c);
    }

    protected void 自动扩容(int index) {
        if (index < 字节码.length) return;
        byte[] 新字节码 = new byte[index * 2];
        System.arraycopy(字节码, 0, 新字节码, 0, 字节码.length);
        字节码 = 新字节码;
    }

    private void 生成操作码(Token id) {
        var 指令名 = id.getText();
        var 操作码 = 指令映射.get(指令名);
        if (操作码 == null) {
            System.err.println("[第" + id.getLine() + "行]指令: " + 指令名 + " 不存在");
            return;
        }

        自动扩容(当前地址 + 1);
        字节码[当前地址] = (byte) (操作码 & 0xFF);
        当前地址++;
    }

    private void 生成操作数(AssemblerParser.OperandContext 操作数) {
        var 文本 = 操作数.getText();
        var token = 操作数.start;

        int 值 = switch (token.getType()) {
            case AssemblerParser.ID -> 取标签地址(文本);
            case AssemblerParser.REG -> 取寄存器号(文本);
            case AssemblerParser.FUNC -> 取函数索引(文本.substring(0, 文本.length() - 2)); // 去掉末尾的 "()"
            case AssemblerParser.INT -> Integer.parseInt(文本);
            case AssemblerParser.CHAR -> 文本.charAt(1); // 去掉单引号
            case AssemblerParser.STRING -> 取常量索引(去括号(文本));
            case AssemblerParser.FLOAT -> 取常量索引(Float.parseFloat(文本));
            default -> {
                System.err.println("[第" + token.getLine() + "行]操作数: " + 文本 + " 类型不支持");
                yield 0;
            }
        };

        自动扩容(当前地址 + 4);
        写整数(字节码, 当前地址, 值);
        当前地址 += 4;
    }

    private int 取标签地址(String 标签名) {
        var 标签符号 = 标签.computeIfAbsent(标签名, 标签符号::new);

        if (标签符号.是否为前向引用) {
            标签符号.添加引用(当前地址);
            return 0; // 暂时还不知道地址
        } else {
            return 标签符号.地址;
        }
    }

    // 转换 "rN" -> N
    private int 取寄存器号(String 寄存器名) {
        return Integer.parseInt(寄存器名.substring(1));
    }

    private int 取函数索引(String 函数名) {
        var 函数符号 = new 函数符号(函数名);
        var 索引 = 常量池.indexOf(函数符号);
        if (索引 >= 0) return 索引;
        return 取常量索引(函数符号);
    }

    private String 去括号(String 文本) {
        return 文本.substring(1, 文本.length() - 1); // 去掉首尾的引号
    }


    public static void 写整数(byte[] 字节码, int 地址, int 值) {
        字节码[地址] = (byte) ((值 >> (8 * 3)) & 0xFF); // get highest byte  // 大端模式，高位在前
        字节码[地址 + 1] = (byte) ((值 >> (8 * 2)) & 0xFF);
        字节码[地址 + 2] = (byte) ((值 >> (8)) & 0xFF);
        字节码[地址 + 3] = (byte) (值 & 0xFF);
    }

    public static int 读整数(byte[] 字节码, int 地址) {
        int b1 = 字节码[地址++] & 0xFF; // mask off sign-extended bits
        int b2 = 字节码[地址++] & 0xFF;
        int b3 = 字节码[地址++] & 0xFF;
        int b4 = 字节码[地址] & 0xFF;
        return b1 << (8 * 3) | b2 << (8 * 2) | b3 << (8) | b4;
    }
    
    public static String 输出字节码(byte[] 字节码, int 字节码尺寸){
        var 结果 = new StringBuffer();
        // 格式：  地址 - 8位
     
        for (int begin = 0; begin < 字节码尺寸; begin += 8) {
            结果.append(String.format("%04d", begin));
            结果.append(":");
            for (int i = 0; i < 8; i++) {
                if (begin + i < 字节码尺寸) {
                    结果.append(String.format("%4d", 字节码[begin + i]));
                } else {
                    break;
                }
            }
            
            结果.append(System.lineSeparator());
        }
        
        return 结果.toString();
    }
}
